import gc
import os
import numpy as np
import cv2
from PIL import Image
from concurrent.futures import ThreadPoolExecutor

Image.MAX_IMAGE_PIXELS = None

def load_large_image(image_path):
    img = Image.open(image_path)
    return np.array(img)

def load_image_tile(image_dir, image_file):
    image_path = os.path.join(image_dir, image_file)
    image = load_large_image(image_path)
    if image is None:
        raise ValueError(f"无法加载图像：{image_path}")
    return image

def load_flow_tile(flow_dir, row, col):
    flow_path = os.path.join(flow_dir, f"tr{row}-tc{col}.npy")
    flow = np.load(flow_path)

    return flow

def apply_flow_to_tile(image, flow):
    h, w = image.shape[:2]
    flow_map = np.zeros((h, w, 2), dtype=np.float32)
    flow_map[..., 0] = flow[..., 0] + np.arange(w)
    flow_map[..., 1] = flow[..., 1] + np.arange(h)[:, np.newaxis]

    remapped_image = cv2.remap(image, flow_map, None, cv2.INTER_CUBIC)
    return remapped_image

def save_tile(image, output_dir, file_name):
    output_path = os.path.join(output_dir, file_name)
    cv2.imwrite(output_path, image)

def process_tile(image_dir, flow_dir, output_dir, image_file):
    # 提取行列编号
    parts = image_file.split('tr')[1].split('-tc')
    row = int(parts[0])
    col = int(parts[1].split('.')[0])

    # 加载当前块及其光流场
    image = load_image_tile(image_dir, image_file)
    flow = load_flow_tile(flow_dir, row, col)

    # 应用光流场
    transformed_tile = apply_flow_to_tile(image, flow)

    # 保存扩展后的图像
    save_tile(transformed_tile, output_dir, image_file)
    # 显式地删除不再需要的变量并调用垃圾收集器
    del image, flow, transformed_tile
    gc.collect()

def process_large_image(image_dir, flow_dir, output_dir, tile_size=1024, border_size=512, max_workers=8):
    os.makedirs(output_dir, exist_ok=True)
    import time

    start_time = time.time()

    image_files = [f for f in os.listdir(image_dir) if f.endswith('.png')]

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(process_tile, image_dir, flow_dir, output_dir, image_file) for image_file in image_files]

        for future in futures:
            future.result()

    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Elapsed time: {elapsed_time} seconds")

# 示例用法
image_dir = '/root/data/deep_align/deep_align/test'  # 替换为实际图像块目录
flow_dir = '/root/data/deep_align/deep_align/test_flow'  # 替换为实际光流场目录
output_dir = '/root/data/deep_align/deep_align/test_result'  # 替换为实际输出目录

process_large_image(image_dir, flow_dir, output_dir)
